# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import copy
import hnswlib
import threading
import numpy as np
import json
import pickle
from utils import get_logger


class HNSWVectorDB:
    def __init__(self, db_name_prefix, max_elements, space="cosine", ef_construction=200, M=16, save_dir=None):
        self.db_name_prefix = db_name_prefix
        self.logger = get_logger(f"HNSWVectorDB: {self.db_name_prefix}")
        self.logger.info(
            f"[HNSWVectorDB] init, max_elements: {max_elements}, space: {space}, ef_construction: {ef_construction}, M: {M}, save_dir: {save_dir}")

        self.save_dir = save_dir

        self.max_elements = max_elements
        self.space = space
        self.ef_construction = ef_construction
        self.M = M

        self._lock = threading.Lock()
        self.vector_db = {}
        self.dynamic_target_ids = set()

    def is_id_exist(self, emb_name, emb_id):
        emb_db = self.vector_db.get(emb_name, None)
        assert emb_db is not None, f"emb_name {emb_name} not exist"
        if not isinstance(emb_id, int):
            try:
                emb_id = int(emb_id)
            except Exception:
                raise TypeError(f"emb_id {emb_id} is not int")
        try:
            items = emb_db.get_items([emb_id])
            return len(items) > 0
        except Exception:
            return False

    def add_data(self, emb_name, emb_id, emb, is_dynamic_target=False):
        assert isinstance(emb, list) and len(emb) > 0, "emb must be a list and len(emb) > 0"
        if not isinstance(emb_id, int):
            try:
                emb_id = int(emb_id)
            except Exception:
                raise TypeError(f"emb_id {emb_id} is not int")

        emb = [float(x) for x in emb]
        with self._lock:
            emb_db = self.vector_db.get(emb_name, None)
            if emb_db is None:
                emb_db = hnswlib.Index(space=self.space, dim=len(emb))
                emb_db.init_index(max_elements=self.max_elements, ef_construction=self.ef_construction, M=self.M)
                self.vector_db[emb_name] = emb_db
            if is_dynamic_target:
                self.dynamic_target_ids.add(emb_id)
        # 判断emb_id是否已经存在了
        is_id_exist = False
        try:
            items = emb_db.get_items([emb_id])
            is_id_exist = len(items) > 0
        except Exception:
            is_id_exist = False
        if is_id_exist:
            return

        emb_db.add_items([emb], [emb_id])

    def search(self, emb_name, query_emb, topk=10, mask_other_dynamic_target=False, keep_dynamic_target_ids=None):
        assert isinstance(query_emb, (list, np.ndarray)) and len(
            query_emb) > 0, "query_emb must be a list and len(query_emb) > 0"
        query_emb = [float(x) for x in query_emb]
        with self._lock:
            emb_db = self.vector_db.get(emb_name, None)
            assert emb_db is not None, f"emb_name {emb_name} not exist"
            emb_db.set_ef(topk)
            mark_deleted_labels = list(self.dynamic_target_ids)
            if mask_other_dynamic_target:
                if keep_dynamic_target_ids is not None and len(keep_dynamic_target_ids) > 0:
                    try:
                        keep_dynamic_target_ids = [int(x) for x in keep_dynamic_target_ids]
                    except Exception:
                        raise TypeError(f"keep_dynamic_target_ids {keep_dynamic_target_ids} is not list of int")
                    mark_deleted_labels = list(self.dynamic_target_ids - set(keep_dynamic_target_ids))
                for mark_id in mark_deleted_labels:
                    emb_db.mark_deleted(mark_id)
            labels, distances = emb_db.knn_query(query_emb, k=topk)
            if mask_other_dynamic_target:
                for mark_id in mark_deleted_labels:
                    emb_db.unmark_deleted(mark_id)
            if len(labels) > 0:
                results = []
                for label, distance in zip(labels[0], distances[0]):
                    results.append((label, 1 - distance))
                return results
            else:
                self.logger.info(f"query search failed")
                return []

    def save_db(self):
        # 将向量库保存到本地
        try:
            db_save_dir = self.save_dir if self.save_dir else "./"
            emb_name_db_path = {}
            with self._lock:
                for emb_name, emb_db in self.vector_db.items():
                    db_name = f"{self.db_name_prefix}-{emb_name}.db.pkl"
                    db_path = os.path.join(db_save_dir, db_name)
                    with open(db_path, "wb") as f:
                        pickle.dump(emb_db, f)
                    emb_name_db_path[emb_name] = db_path
            emb_name_db_path_save_path = os.path.join(db_save_dir, f"{self.db_name_prefix}-emb_name_db_path.json")
            with open(emb_name_db_path_save_path, "w") as f:
                json.dump(emb_name_db_path, f, ensure_ascii=False)
            # 记录的动态targetID集合也需要dump
            dynamic_target_ids_save_path = os.path.join(db_save_dir, f"{self.db_name_prefix}-dynamic_target_ids.pkl")
            with open(dynamic_target_ids_save_path, "wb") as f:
                pickle.dump(self.dynamic_target_ids, f)
            self.logger.info(f"save vector db to local successfully")
        except Exception as e:
            self.logger.error(f"save vector db to local failed, error: {e}")

    def try_load_db(self):
        # 尝试从本地加载向量库
        self.logger.info("try load vector db from local")
        db_save_dir = self.save_dir if self.save_dir else "./"
        emb_name_db_path_save_path = os.path.join(db_save_dir, f"{self.db_name_prefix}-emb_name_db_path.json")
        if not os.path.exists(emb_name_db_path_save_path):
            self.logger.info("local vector db not exists, init new vector db")
            return False
        try:
            with open(emb_name_db_path_save_path, "r") as f:
                emb_name_db_path = json.load(f)
            with self._lock:
                for emb_name, db_path in emb_name_db_path.items():
                    if not os.path.exists(db_path):
                        self.logger.info(f"local vector db {db_path} not exists, init new vector db")
                        return False
                    with open(db_path, "rb") as f:
                        emb_db = pickle.load(f)
                    self.vector_db[emb_name] = emb_db
            # 恢复动态targetID集合
            dynamic_target_ids_save_path = os.path.join(db_save_dir, f"{self.db_name_prefix}-dynamic_target_ids.pkl")
            if not os.path.exists(dynamic_target_ids_save_path):
                self.logger.info("local dynamic target ids not exists, init new dynamic target ids")
                return False
            with open(dynamic_target_ids_save_path, "rb") as f:
                self.dynamic_target_ids = pickle.load(f)
            self.logger.info(f"load vector db from local successfully")
            return True
        except Exception as e:
            self.logger.error(f"load vector db from local failed, error: {e} ")
            return False
